import numpy
import Genefinder
from math import log
from numpy import loadtxt, savetxt
from Genefinder import pi, A
import parser


def viterbi(genome):
    AMatrix = loadtxt("A.txt", delimiter=',')
    pMatrix = loadtxt("p.txt", delimiter=',')
    genome = parser.Parser(genome)
    
    omega = numpy.arange(16.0*genome.length).reshape(16, genome.length)
    for e in omega.flat:
        omega.itemset(e,0)
    symbol = genome.nextsymbol()

    #Base case
    for i in range(16):
        if A[0,i]!=0:
            omega[i,0] = pi[i,0]*pMatrix[Genefinder.mapsymbol(symbol), i]
            
        print[omega[i,0]]

    #Recursion case:
    #prob = 0
    for i in range(1,genome.length):
        for j in range(16):
            prob = 0
            for z in range(16):
                tmp = AMatrix[j,z]*omega[z, i-1]*pMatrix[Genefinder.mapsymbol(symbol),j]
                if tmp > prob:
                    prob = tmp
            omega[j,i] = prob
        symbol = genome.nextsymbol()
            
    savetxt("omega.txt", omega, '%f', ',') 
    
def backtrace(genome):
    genome = parser.Parser(genome)
    genome.gotoend()
    symbol = genome.peek(0)

    print "loading files"
    AMatrix = loadtxt("A.txt", delimiter=',')
    print "A loaded"
    pMatrix = loadtxt("p.txt", delimiter=',')
    print "p loaded"
    omega = loadtxt("omega.txt", delimiter=',')
    print "loading files completed"

    arg_k = 0
    max = 0
    for k in range(16):
        possibleMax = omega[k, genome.length-1]
        if possibleMax > max:
            max = possibleMax
            arg_k = k            
    resultFile = [arg_k]
    print "set last result to:"
    print arg_k
    
    for n in reversed(range(genome.length-1)):
        max = 0.0
        arg_k = 0
        for k in range(16):
            n_1 = resultFile[genome.length-2-n]
            possibleMax = pMatrix[Genefinder.mapsymbol(symbol),n_1]*omega[k,n]*AMatrix[k,n_1]
            if possibleMax > max:
                max = possibleMax
                arg_k = k
        print arg_k
        print max
        resultFile.append(arg_k)
        symbol = genome.prevsymbol()
    
    with open("intermediate.txt", "w") as f:
        f.write("".join([str(s) for s in resultFile]))
    
#Doesn't work.. need to convert to hex or something first.. to differenciate between 1 and 15 fx
def stateToFasta(fileName):
    states = parser.Parser(fileName)
    state = states.nextsymbol()
    with open("result.fa", "w") as f:
        while state:
            if state == 0:
                f.write("N")
            else:
                f.write("C")
    

viterbi("genome9Copy.fa")
backtrace("genome9Copy.fa")
